import argparse
from database import DatabaseManager
from config import Config

def main():
    parser = argparse.ArgumentParser(description="Database Sensitive Data Replacement Tool")
    parser.add_argument("table_name", help="Name of the table to process")
    parser.add_argument("sensitive_columns", nargs='+', help="Sensitive columns to replace")
    parser.add_argument("mapping_table_name", help="Name of the mapping table to create")  # 新增参数：映射表名称
    args = parser.parse_args()

    config = Config()
    db_manager = DatabaseManager(config)

    table_name = args.table_name
    sensitive_columns = args.sensitive_columns
    mapping_table_name = args.mapping_table_name  # 获取用户输入的映射表名称

    # Replace sensitive data in original table
    db_manager.replace_sensitive_data(table_name, sensitive_columns, mapping_table_name)  # 传递映射表名称

    print(f"Sensitive data replacement completed for table {table_name}")

if __name__ == "__main__":
    main()